package org.budo.java.agent.transformer.javassist;

import java.io.ByteArrayInputStream;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.reflect.Modifier;
import java.security.ProtectionDomain;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;

import org.budo.java.agent.AgentInvoker;
import org.budo.java.agent.enhancer.AgentEnhancer;
import org.budo.java.agent.filter.RuntimeAgentFilterChain;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.CtNewMethod;
import javassist.bytecode.AnnotationsAttribute;
import javassist.bytecode.MethodInfo;

/**
 * @author lmw
 */
public class JavassistClassFileTransformer extends AbstractJavassistClassFileTransformer {
    public JavassistClassFileTransformer() {
    }

    public JavassistClassFileTransformer(String agentArgs) {
        this.setAgentArgs(agentArgs);
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classFileBuffer) throws IllegalClassFormatException {
        System.err.println("#35 className=" + className);

        if (null == className) {
            return null;
        }

        if (className.startsWith("java/") || className.startsWith("javax/") || className.startsWith("jdk/") || className.startsWith("sun/")) {
            return null;
        }

        if (className.startsWith("org/eclipse/jdt/")) {
            return null;
        }

        if (className.startsWith("org/budo/java/agent/")) {
            return null;
        }

        try {
            ClassPool classPool = ClassPool.getDefault();
            CtClass ctClass = classPool.makeClass(new ByteArrayInputStream(classFileBuffer));

            if (ctClass.isInterface()) { // 接口跳过
                return null;
            }

            if (Modifier.isFinal(ctClass.getModifiers())) {
                return null; // final 类跳过
            }

            ctClass = this.transform_1(ctClass);
            ctClass.detach(); // 将内存中曾经被javassist加载过的Date对象移除

            return ctClass.toBytecode();
        } catch (Throwable e) {
            throw new RuntimeException(e);
        }
    }

    public CtClass transform_1(CtClass ctClass) {
        try {
            ClassPool classPool = ClassPool.getDefault();

            ctClass.addInterface(classPool.get(AgentInvoker.class.getName())); // 加接口

            Map<Integer, CtMethod> methodMap = new HashMap<Integer, CtMethod>();

            CtMethod[] methods = ctClass.getMethods();

            String[] methodNames = new String[methods.length];
            CtClass[][] methodParameterTypes = new CtClass[methods.length][];

            for (int i = 0; i < methods.length; i++) {
                CtMethod method = methods[i];

                if (Modifier.isFinal(method.getModifiers())) {
                    continue; // final 方法跳过
                }

                if (!method.getDeclaringClass().equals(ctClass)) {
                    continue; // 非本类声明的方法跳过
                }

                String originalMethodName = method.getName();
                method.setName("$" + originalMethodName); // 原方法重命名

                MethodInfo methodInfo = method.getMethodInfo();
                AnnotationsAttribute annotationsAttribute = (AnnotationsAttribute) methodInfo.getAttribute(AnnotationsAttribute.visibleTag); // 取得注解，加到替代方法
                methodInfo.removeAttribute(AnnotationsAttribute.visibleTag); // 重命名的原方法去掉注解

                CtMethod delegator = this.buildMethodDelegator(method, originalMethodName, i, annotationsAttribute); // 添加原方法的替代
                ctClass.addMethod(delegator);

                methodNames[i] = originalMethodName;
                methodParameterTypes[i] = method.getParameterTypes();
                methodMap.put(i, method);
            }

            this.putMethodName(ctClass, methodNames);
            this.putMethodParameterTypes(ctClass, methodParameterTypes);

            ctClass.addMethod(this.buildInvokerMethod(ctClass, methodMap));

            return ctClass;
        } catch (Throwable e) {
            throw new RuntimeException(e);
        }
    }

    private String buildMethodDelegatorBody(CtMethod method, Integer methodIndex) throws Throwable {
        String callChain = AgentEnhancer.class.getName() + ".buildChain(" + method.getDeclaringClass().getName() + ".class, " + methodIndex + ", this, $args).doFilter().getResult()";
        CtClass returnType = method.getReturnType();
        String returnTypeName = returnType.getName();
        if (returnTypeName.equals("void")) {
            return callChain + "; ";
        }

        return "return " + this.toPrimitiveValue(returnType, callChain) + "; ";
    }

    private CtMethod buildMethodDelegator(CtMethod method, String originalMethodName, Integer methodIndex, AnnotationsAttribute annotationsAttribute) throws Throwable {
        String methodBody = this.buildMethodDelegatorBody(method, methodIndex);

        CtMethod delegator = CtNewMethod.make(method.getModifiers(), //
                method.getReturnType(), //
                originalMethodName, //
                method.getParameterTypes(), //
                method.getExceptionTypes(), //
                methodBody, //
                method.getDeclaringClass());

        if (null != annotationsAttribute) {
            delegator.getMethodInfo().addAttribute(annotationsAttribute);
        }

        return delegator;
    }

    private CtMethod buildInvokerMethod(CtClass ctClass, Map<Integer, CtMethod> methodMap) throws Throwable {
        String src = " public void invoke(" + RuntimeAgentFilterChain.class.getName() + " filterChain) { ";
        for (Entry<Integer, CtMethod> entry : methodMap.entrySet()) {
            src += " if ( " + entry.getKey() + " == filterChain.getMethodIndex() ) { ";
            CtMethod method = entry.getValue();

            String callMethod = method.getName() + "(";
            CtClass[] parameterTypes = method.getParameterTypes();
            if (parameterTypes.length > 0) {
                callMethod += this.toPrimitiveValue(parameterTypes[0], "filterChain.getArgs()[" + 0 + "]");
            }
            for (int i = 1; i < parameterTypes.length; i++) {
                callMethod += ", " + this.toPrimitiveValue(parameterTypes[1], "filterChain.getArgs()[" + 1 + "]");
            }
            callMethod += " ) ";

            CtClass returnType = method.getReturnType();
            String returnTypeName = returnType.getName();
            if (returnTypeName.equals("void")) {
                src += callMethod + "; return; ";
            } else {
                src += "filterChain.setResult( " + this.toObjectValue(returnType, callMethod) + " ) ; return; ";
            }
            src += " } ";
        }
        src += " } ";

        return CtMethod.make(src, ctClass);
    }
}